package com.yntx.usercenter.common.config.sharding;

import cn.hutool.core.collection.CollectionUtil;
import cn.hutool.extra.spring.SpringUtil;
import com.alibaba.druid.util.StringUtils;
import com.yntx.usercenter.common.config.sharding.enums.ShardingTableCacheEnum;
import lombok.extern.slf4j.Slf4j;
import org.apache.shardingsphere.driver.jdbc.core.datasource.ShardingSphereDataSource;
import org.apache.shardingsphere.infra.config.RuleConfiguration;
import org.apache.shardingsphere.infra.database.DefaultSchema;
import org.apache.shardingsphere.infra.hint.HintManager;
import org.apache.shardingsphere.mode.manager.ContextManager;
import org.apache.shardingsphere.sharding.algorithm.config.AlgorithmProvidedShardingRuleConfiguration;
import org.apache.shardingsphere.sharding.algorithm.sharding.ShardingAlgorithmException;
import org.apache.shardingsphere.sharding.api.config.rule.ShardingTableRuleConfiguration;
import org.springframework.core.env.Environment;

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.*;
import java.util.stream.Collectors;

/**
 * <p> @Title ShardingAlgorithmTool
 * <p> @Description 按月分片算法工具
 *
 * @author ACGkaka
 * @date 2022/12/20 14:03
 */
@Slf4j
public class ShardingAlgorithmTool {

    /**
     * 表分片符号，例：t_contract_202201 中，分片符号为 "_"
     */
    private static final String TABLE_SPLIT_SYMBOL = "_";

    /**
     * 数据库配置
     */
    private static final Environment ENV = SpringUtil.getApplicationContext().getEnvironment();

    /**
     * 检查分表获取的表名是否存在，不存在则自动建表
     *
     * @param logicTable       逻辑表
     * @param resultTableNames 真实表名，例：t_contract_202201
     * @return 存在于数据库中的真实表名集合
     */
    public static Set<String> getShardingTablesAndCreate(ShardingTableCacheEnum logicTable, Collection<String> resultTableNames) {
        return resultTableNames.stream().map(o -> getShardingTableAndCreate(logicTable, o)).collect(Collectors.toSet());
    }

    /**
     * 分片表和创建
     * 检查分表获取的表名是否存在，不存在则自动建表
     *
     * @param logicTable      逻辑表
     * @param resultTableName 真实表名，例：t_contract_202201
     * @return 确认存在于数据库中的真实表名
     */
    public static String getShardingTableAndCreate(ShardingTableCacheEnum logicTable, String resultTableName) {
        Collection<Comparable<?>> databaseShardingValues = HintManager.getDatabaseShardingValues(logicTable.logicTableName());
        List<String> dbName = databaseShardingValues.stream().map(Objects::toString).collect(Collectors.toList());
        if (CollectionUtil.isEmpty(dbName)) {
            throw new ShardingAlgorithmException("找不到上下文分库库名");
        }
        // 缓存中有此表则返回，没有则判断创建
        if (logicTable.resultTableNamesCache().contains(dbName.get(0) + "." + resultTableName)) {
            return resultTableName;
        } else {
            // 未创建的表返回逻辑空表
            boolean isSuccess = createShardingTable(logicTable, dbName.get(0), resultTableName);
//            return isSuccess ? resultTableName : logicTable.logicTableName();
            return resultTableName;
        }
    }

    /**
     * 重载全部缓存
     */
    public static void tableNameCacheReloadAll() {
        Arrays.stream(ShardingTableCacheEnum.values()).forEach(ShardingAlgorithmTool::tableNameCacheReload);
    }

    /**
     * 重载指定分表缓存
     *
     * @param shardingTableCacheEnum 逻辑表，例：t_contract
     */
    public static void tableNameCacheReload(ShardingTableCacheEnum shardingTableCacheEnum) {
        // 读取数据库中|所有表名
        List<String> tableNameList = getAllActualTableName(shardingTableCacheEnum);
        // 删除旧的缓存（如果存在）
        shardingTableCacheEnum.resultTableNamesCache().clear();
        // 写入新的缓存
        shardingTableCacheEnum.resultTableNamesCache().addAll(tableNameList);
        // 动态更新配置 actualDataNodes
        actualDataNodesRefresh(shardingTableCacheEnum);
    }

    /**
     * 获取所有实际表名
     *
     * @param logicTable 逻辑表
     * @return 表名集合
     */
    public static List<String> getAllActualTableName(ShardingTableCacheEnum logicTable) {
        List<String> tableNames = new ArrayList<>();
        ShardingSphereDataSource shardingSphereDataSource = SpringUtil.getBean(ShardingSphereDataSource.class);
        Map<String, DataSource> dataSourceMap = shardingSphereDataSource.getContextManager().getDataSourceMap(DefaultSchema.LOGIC_NAME);
        dataSourceMap.forEach((dbName, dataSource) -> {
            try (Connection conn = dataSource.getConnection()) {
                Statement st = conn.createStatement();
                String logicTableName = logicTable.logicTableName();
                // 兼容不分表的情况，逻辑表实际上也对应物理存在
                tableNames.add(dbName + "." + logicTableName);
                try (ResultSet rs = st.executeQuery("show TABLES like '" + logicTableName + TABLE_SPLIT_SYMBOL + "%'")) {
                    while (rs.next()) {
                        String tableName = rs.getString(1);
                        // 匹配分表格式 例：^(t\_user_\d)$
                        if (tableName != null && tableName.matches(String.format("^(%s\\d+)$", logicTableName + TABLE_SPLIT_SYMBOL))) {
                            tableNames.add(dbName + "." + rs.getString(1));
                        }
                    }
                }
            } catch (SQLException e) {
                log.error("#getAllTableNameBySchema# connect database error:{}", e.getMessage(), e);
                throw new IllegalArgumentException("数据库连接失败");
            }
        });
        return tableNames;
    }

    /**
     * 动态更新配置 actualDataNodes
     *
     * @param logicTable
     */
    public static void actualDataNodesRefresh(ShardingTableCacheEnum logicTable) {
        try {
            // 获取数据分片节点
            String logicTableName = logicTable.logicTableName();
            Set<String> tableNamesCache = logicTable.resultTableNamesCache();
            log.info("#actualDataNodesRefresh# logicTableName:{}，tableNamesCache:{}", logicTableName, tableNamesCache);

            // generate actualDataNodes
            String newActualDataNodes = String.join(",", tableNamesCache);
            ShardingSphereDataSource shardingSphereDataSource = SpringUtil.getBean(ShardingSphereDataSource.class);
            updateShardRuleActualDataNodes(shardingSphereDataSource, logicTableName, newActualDataNodes);
        } catch (Exception e) {
            log.error("#actualDataNodesRefresh# error:{}", e.getMessage(), e);
        }
    }

    /**
     * 刷新ActualDataNodes
     */
    private static void updateShardRuleActualDataNodes(ShardingSphereDataSource dataSource, String logicTableName, String newActualDataNodes) {
        if (StringUtils.isEmpty(newActualDataNodes)) {
            log.warn("#updateShardRuleActualDataNodes# error: newActualDataNodes is empty!");
            return;
        }
        // Context manager.
        ContextManager contextManager = dataSource.getContextManager();

        // Rule configuration.
        String schemaName = "logic_db";
        Collection<RuleConfiguration> newRuleConfigList = new LinkedList<>();
        Collection<RuleConfiguration> oldRuleConfigList = dataSource.getContextManager()
                .getMetaDataContexts()
                .getMetaData(schemaName)
                .getRuleMetaData()
                .getConfigurations();

        for (RuleConfiguration oldRuleConfig : oldRuleConfigList) {
            if (oldRuleConfig instanceof AlgorithmProvidedShardingRuleConfiguration) {

                // Algorithm provided sharding rule configuration
                AlgorithmProvidedShardingRuleConfiguration oldAlgorithmConfig = (AlgorithmProvidedShardingRuleConfiguration) oldRuleConfig;
                AlgorithmProvidedShardingRuleConfiguration newAlgorithmConfig = new AlgorithmProvidedShardingRuleConfiguration();

                // Sharding table rule configuration Collection
                Collection<ShardingTableRuleConfiguration> newTableRuleConfigList = new LinkedList<>();
                Collection<ShardingTableRuleConfiguration> oldTableRuleConfigList = oldAlgorithmConfig.getTables();

                oldTableRuleConfigList.forEach(oldTableRuleConfig -> {
                    if (logicTableName.equals(oldTableRuleConfig.getLogicTable())) {
                        ShardingTableRuleConfiguration newTableRuleConfig = new ShardingTableRuleConfiguration(oldTableRuleConfig.getLogicTable(), newActualDataNodes);
                        newTableRuleConfig.setTableShardingStrategy(oldTableRuleConfig.getTableShardingStrategy());
                        newTableRuleConfig.setDatabaseShardingStrategy(oldTableRuleConfig.getDatabaseShardingStrategy());
                        newTableRuleConfig.setKeyGenerateStrategy(oldTableRuleConfig.getKeyGenerateStrategy());

                        newTableRuleConfigList.add(newTableRuleConfig);
                    } else {
                        newTableRuleConfigList.add(oldTableRuleConfig);
                    }
                });

                newAlgorithmConfig.setTables(newTableRuleConfigList);
                newAlgorithmConfig.setAutoTables(oldAlgorithmConfig.getAutoTables());
                newAlgorithmConfig.setBindingTableGroups(oldAlgorithmConfig.getBindingTableGroups());
                newAlgorithmConfig.setBroadcastTables(oldAlgorithmConfig.getBroadcastTables());
                newAlgorithmConfig.setDefaultDatabaseShardingStrategy(oldAlgorithmConfig.getDefaultDatabaseShardingStrategy());
                newAlgorithmConfig.setDefaultTableShardingStrategy(oldAlgorithmConfig.getDefaultTableShardingStrategy());
                newAlgorithmConfig.setDefaultKeyGenerateStrategy(oldAlgorithmConfig.getDefaultKeyGenerateStrategy());
                newAlgorithmConfig.setDefaultShardingColumn(oldAlgorithmConfig.getDefaultShardingColumn());
                newAlgorithmConfig.setShardingAlgorithms(oldAlgorithmConfig.getShardingAlgorithms());
                newAlgorithmConfig.setKeyGenerators(oldAlgorithmConfig.getKeyGenerators());

                newRuleConfigList.add(newAlgorithmConfig);
            }
        }

        // update context
        contextManager.alterRuleConfiguration(schemaName, newRuleConfigList);
    }

    /**
     * 创建分表
     *
     * @param logicTable      逻辑表
     * @param resultTableName 真实表名，例：t_contract_202201
     * @return 创建结果（true创建成功，false未创建）
     */
    private static boolean createShardingTable(ShardingTableCacheEnum logicTable, String dbName, String resultTableName) {
        synchronized (logicTable.logicTableName().intern()) {
            // 缓存中有此表 返回
            if (logicTable.resultTableNamesCache().contains(dbName + "." + resultTableName)) {
                return false;
            }
            // 缓存中无此表，则建表并添加缓存
            createTable(Collections.singletonList("CREATE TABLE IF NOT EXISTS `" + resultTableName + "` LIKE `" + logicTable.logicTableName() + "`;"));
            // 缓存重载
            tableNameCacheReload(logicTable);
        }
        return true;
    }

    /**
     * 执行SQL
     *
     * @param sqlList SQL集合
     */
    private static void createTable(List<String> sqlList) {
        // todo：替换为传入dbName的mapper操作
//        if (tenant == null) {
//            throw new IllegalArgumentException("没有" + tenant + "数据源");
//        }
//        String url = tenant.getDatasourceUrl();
//        String username = tenant.getDatasourceUsername();
//        String password = tenant.getDatasourcePassword();
//        if (StringUtils.isEmpty(url) || StringUtils.isEmpty(username) || StringUtils.isEmpty(password)) {
//            log.error(">>>>>>>>>> 【ERROR】数据库连接配置有误，请稍后重试，URL:{}, username:{}, password:{}", url, username, password);
//            throw new IllegalArgumentException("数据库连接配置有误，请稍后重试");
//        }
//        try (Connection conn = DriverManager.getConnection(url, username, password)) {

        ShardingSphereDataSource shardingSphereDataSource = SpringUtil.getBean(ShardingSphereDataSource.class);
        try (Connection conn = shardingSphereDataSource.getConnection()) {
            try (Statement st = conn.createStatement()) {
                conn.setAutoCommit(false);
                for (String sql : sqlList) {
                    st.execute(sql);
                }
                Thread.sleep(500L);
            } catch (Exception e) {
                conn.rollback();
                log.error("#createTable# execute create sql error:{}", e.getMessage(), e);
                throw new IllegalArgumentException("数据表创建执行失败");
            }
        } catch (SQLException e) {
            log.error("#createTable# connect database error:{}", e.getMessage(), e);
            throw new IllegalArgumentException("数据库连接失败");
        }
    }

}

